import sys
import os
import glob
import math
import subprocess
import tarfile
import datetime
import shutil
from netCDF4 import Dataset
from PIL import Image, ImageDraw
import pytz

import common
import dataStats

radarTestDir = "%s/data/ktlx-radar/2010/KTLX20100301"%common.BaseDir

def getClosestRadarData(time, bNum = 1, fNum = 1, product='baseRef0'):
	"""bNum = number of frames before the time; 
	   fNum = number of frames after the time"""
	productsDict = {"baseRef0":"N0RTLX", "baseRef1":"N1RTLX",
				"baseRef2":"N2RTLX", "baseRef3":"N3RTLX",
				"stormRel0":"N0STLX", "stormRel1":"N1STLX",
				"stormRel2":"N2STLX", "stormRel3":"N3STLX",
				"baseVel0":"N0VTLX", "baseVel1":"N1VTLX",
				"baseVel2":"N2VTLX", "baseVel3":"N3VTLX"}
	
	curDay = datetime.datetime(time.year, time.month, time.day)
	prevDay = curDay + datetime.timedelta(days = -1)
	
	path1 = "%s/%d/KTLX%d%.2d%.2d"%(common.RadarBaseDir, curDay.year, curDay.year, curDay.month, curDay.day)
	path2 = "%s/%d/KTLX%d%.2d%.2d"%(common.RadarBaseDir, prevDay.year, prevDay.year, prevDay.month, prevDay.day)
	fileList1 = glob.glob("%s/*%s*"%(path1, productsDict[product]))
	fileList2 = glob.glob("%s/*%s*"%(path2, productsDict[product]))
	filteredFileList = [f for f in (fileList2+fileList1) if len(f.split('.'))==1]
	
	utc = pytz.utc
	
	closestIndex = -1
	minTimeDeltaSecs = -1
	for i in range(len(filteredFileList)):
		f = filteredFileList[i]
		fParts = f.split('\\')
		t = fParts[-1][-4:]
		
		hour = int(t[0:2])
		minute = int(t[2:4])
		curTime = utc.localize(datetime.datetime(time.year, time.month, time.day,
									hour, minute, 0))
	
		if abs((curTime - time).total_seconds()) < minTimeDeltaSecs or minTimeDeltaSecs == -1:
			closestIndex = i
			minTimeDeltaSecs = abs((curTime - time).total_seconds())
	
	begin = closestIndex-bNum
	end = closestIndex+fNum+1
	if begin < 0: begin = 0
	if end >= len(filteredFileList): end = len(filteredFileList)-1
	
	return filteredFileList[begin:end]
    
def radarToNetCDF(radarFiles, recursive = False, baseDir = common.RadarBaseDir, verbose = True):
    ncFiles = []

    if recursive:    
        for prefix in common.RadarPrefixes:
            filelist = glob.glob("%s/KOUN*%s*"%(baseDir, prefix))
    else:
        filelist = radarFiles

    for f in filelist:
        fParts = f.split('.')
        if len(fParts) == 1 and not os.path.exists("%s.nc"%f):
            subprocess.check_call("java -classpath %s/toolsUI-4.2.jar ucar.nc2.FileWriter -in %s -out %s.nc"%(common.RadarAnalysisDir, f, f), shell=True)
            ncFiles.append("%s.nc"%f)
            if verbose:
                print "processed %s"%f
        elif os.path.exists("%s.nc"%f):
            ncFiles.append("%s.nc"%f)
    return ncFiles

def netCDFToImg(ncFiles, prefix = "", verbose = True):
    for ncFileName in ncFiles:

        flatData, curPrefix = getRadarData(ncFileName, prefix)
        curRadarVarName = common.RadarVariables[curPrefix]

        fParts = ncFileName.split('.')

        if curRadarVarName == 'StormMeanVelocity':
            im = Image.new('RGB', (230, 360))

            def buildRGB(v):
                if v > 0:
                    return (int(v), 0, 0)
                else:
                    return (0, -int(v), 0)
            
            im.putdata([buildRGB(x) for x in flatData])
            im.save("%s.bmp"%fParts[0], 'bmp')
        else:
            im = Image.new('L', (230, 360))
            im.putdata(flatData)
            im.save("%s.bmp"%fParts[0], 'bmp')
        if verbose:
            print "saved image %s.bmp"%fParts[0]

def radarToImg(radarData, prefix, outName, center = None, dim = None, verbose = True, mark = None, markLen = 5):
    """radarData is a 360x230 array (so len(radarData) = 360, len(radarData[0]) = 230)"""

    curRadarVarName = common.RadarVariables[prefix]

    if dim == None:
        dim = (360, 230)

    if center == None:
        center = (len(radarData)/2, len(radarData[0])/2)

    lb = center[0] - dim[0]/2
    rb = center[0] + dim[0]/2
    tb = center[1] - dim[1]/2
    bb = center[1] + dim[1]/2

    if mark == (-1, -1):
        subX = 0
        subY = 0
        if lb < 0:
            subX = lb
        if tb < 0:
            subY = tb
        mark = (subX+(rb-lb)/2.0, subY+(bb-tb)/2.0)

    if lb < 0: lb = 0
    if rb >= len(radarData): rb = len(radarData)
    if tb < 0: tb = 0
    if bb >= len(radarData[0]): bb = len(radarData[0])

    im = Image.new('RGB', ((rb-lb), (bb-tb)))
    flatData = [radarData[angle][dist] for dist in range(tb, bb) for angle in range(lb, rb)]
    
    if curRadarVarName == 'StormMeanVelocity':
        def buildRGB(v):
            if v > 0:
                return (int(v)*4, 0, 0)
            else:
                return (0, -int(v)*4, 0)

        im.putdata([buildRGB(x) for x in flatData])
    else:
        def buildRGB(v):
            return common.RadarColors[int(5*(round(float(v)/5)))]
            
        im.putdata([buildRGB(x) for x in flatData])

    # center
    if mark != None:
        draw = ImageDraw.Draw(im)

        xLeft = markLen
        if mark[0]-xLeft < 0: xLeft = mark[0]

        xRight = markLen
        if mark[0]+xRight >= dim[0]: xRight = (dim[0]-mark[0])

        yTop = markLen
        if mark[1]-yTop < 0: yTop = mark[1]

        yBottom = markLen
        if mark[1]+yBottom >= dim[1]: yBottom = (dim[1]-mark[1])

        if markLen == 0:
            draw.point((mark[0], mark[1]), fill = (255, 255, 255))
        else:   
            draw.line((mark[0]-xLeft, mark[1], mark[0]+xRight, mark[1]), fill = (255, 255, 255))
            draw.line((mark[0], mark[1]-yTop, mark[0], mark[1]+yBottom), fill = (255, 255, 255))

    im.save("%s.bmp"%outName, 'bmp')
    if verbose:
        print "saved image %s.bmp"%outName
    
def getRadarData(ncFileName, prefix = "", flat = True, flip = True):
    """Takes in an NETCDF file and returns the main radar data from it."""
    curPrefix = prefix
    curRadarVarName = ''
    # determine prefix from file name
    if prefix == "":
        for p in common.RadarPrefixes:
            if ncFileName.find(p) != -1:
                curPrefix = p
                curRadarVarName = common.RadarVariables[curPrefix]
                
    else:
        curRadarVarName = common.RadarVariables[prefix]
                
    # failed to find prefix
    if curRadarVarName == "":
        return ([], '')

    fParts = ncFileName.split('.')
    ncFile = Dataset(ncFileName, formst='NETCDF3')
    ncFileRef = ncFile.variables[curRadarVarName][:]
    azimuthData = ncFile.variables['azimuth'][:]

    radarData = []

    # make the correction for angle
    for ray in range(len(ncFileRef)):
        radarData.append((azimuthData[ray], ncFileRef[ray]))

    radarData.sort(key=(lambda x:x[0]))

    def transformRadarData(v):
        if math.isnan(v): return 0
        else: return v

    retData = []
    if flat:
        retData = [transformRadarData(x) for tup in radarData for x in tup[1]]
    else:
        if flip:
            for i in range(len(radarData)):
                tup = radarData[i]
                newRow = []
                for j in range(len(tup[1])):
                    newRow.append(transformRadarData(tup[1][len(tup[1])-1-j]))
                retData.append(newRow)

        else:
            for i in range(len(radarData)):
                tup = radarData[i]
                newRow = []
                for j in range(len(tup[1])):
                    newRow.append(transformRadarData(tup[1][j]))
                retData.append(newRow)
            
    ncFile.close()
    return (retData, curPrefix)        

def cleanRadarData(baseDir = common.RadarBaseDir, deleteExtensions = ['nc', 'bmp'], verbose = True):
    """Removes all radar files with prefixes other than those passed in in the prefixes list."""
    
    for dirname, dirnames, filenames in os.walk(baseDir):
        cleaned = False
        if verbose:
            print dirname
        for f in glob.glob("%s/KOUN*"%dirname):
            fParts = f.split('.')
            ext = fParts[-1]

            if len(fParts) > 1:
                if ext in deleteExtensions:
                    os.remove("%s"%f)
                    cleaned = True
            else:
                foundPrefix = False
                for prefix in common.RadarPrefixes:
                    if f.find(prefix) != -1:
                        foundPrefix = True
                        break
                if not foundPrefix:
                    os.remove("%s"%f)
                    cleaned = True
        if cleaned and verbose:
            print "cleaned %s"%(dirname)


def unzipRadarData(years = [], baseDir = common.RadarBaseDir, overwrite = False, clean = True, verbose = True):
	"""First unzips the radar data, then untars the resulting tar file."""
	for dirname, dirnames, filenames in os.walk(baseDir):
		examine = False
		for year in years:
			if dirname.find(str(year)) != -1:
				examine = True
		if not examine:
			continue
		
		print 'examining %s'%dirname
		if len(filenames) == 1 or overwrite:
			fList = glob.glob("%s/*.tar.Z"%dirname)
			# no .tar.Z file
			if len(fList) == 0:
				# look for .tar files...
				fList = glob.glob("%s/*.tar"%dirname)
                
				if len(fList) == 0:
					continue
	
			fParts = fList[0].split('.')
			if fList[0].find(".tar.Z") != -1:
				subprocess.check_call("%s/gzip.exe -d %s"%(common.GzipDir, os.path.join(dirname, filenames[0])), shell=True)
			if os.path.exists(os.path.join(dirname, "%s.tar"%fParts[0])):
				tar = tarfile.TarFile(os.path.join(dirname, "%s.tar"%fParts[0]))
				tar.extractall(dirname)
				tar.close()
				os.remove(os.path.join(dirname, "%s.tar"%fParts[0]))
				if clean:
					cleanRadarData(dirname)
				if verbose:
					print "untarred %s"%os.path.join(dirname, "%s.tar"%fParts[0])

def organizeRadarTars(years, baseDir = common.RadarBaseDir, verbose = True):
    """Takes a group of .tar.Z radar files and puts each in its own folder."""    
    for year in years:
        filelist = glob.glob("%s/%s/*.Z"%(baseDir, year))
        for fname in filelist:
            fnameparts = fname.split("\\")
            fnameparts_no_ext = fnameparts[1].split('.')
            if not os.path.exists('%s/%s'%(fnameparts[0],fnameparts_no_ext[0])):
                os.mkdir('%s/%s'%(fnameparts[0],fnameparts_no_ext[0]))
            if not os.path.exists('%s/%s/%s'%(fnameparts[0],fnameparts_no_ext[0],fnameparts[1])):
                if verbose:
                    print 'move %s to %s/%s/%s'%(fname, fnameparts[0], fnameparts_no_ext[0], fnameparts[1])
                shutil.move(fname, '%s/%s/%s'%(fnameparts[0], fnameparts_no_ext[0], fnameparts[1]))

def radarToRv3(radarFolders, deleteTar = False):
    """Extracts radar data from the .tar.Z file for the specificied day using
       the extraction program provided by GrLevelX."""
    for radarFolder in radarFolders:
        if os.path.exists(radarFolder):
            radarName = radarFolder.split('/')[-1]            
            if os.path.exists("%s/%s.tar.Z"%(radarFolder, radarName)):
                if not os.path.exists('%s/rv3'%radarFolder):
                    os.mkdir('%s/rv3'%radarFolder)
                os.chdir("%s/rv3"%radarFolder)
                subprocess.check_call(["%s/ncdctorv3.exe"%(common.RadarBaseDir), "%s/%s.tar.Z"%(radarFolder, radarName)])
                if deleteTar:
                    os.remove("%s/%s.tar.Z"%(radarFolder, radarName))
                

def removeInactiveRadarData(years, baseDir = common.RadarBaseDir):
	"""Removes all radar data for days with no tornado activity."""
	activeDays = dataStats.findRequiredRadarDays(years = years)
	
	activeDateTime = []
	activeFolders = []
	for activeDay in activeDays:
		activeTime = datetime.datetime.strptime(activeDay, '%Y-%m-%d %H:%M:%S')
		activeDateTime.append(activeTime)
		activeFolders.append('%s/%d/KTLX%d%s%s'%(baseDir, activeTime.year, activeTime.year, 
												 str(activeTime.month).zfill(2), 
												 str(activeTime.day).zfill(2)))
	
	for year in years:
		curRadarDir = '%s/%d'%(baseDir, year)
		for dirname, dirnames, filenames in os.walk(curRadarDir):
			for curDirName in dirnames:			
				dirNameParts = curDirName.split('\\')
				curDayStr = dirNameParts[-1]
				
				if curDayStr != 'rv3':
					curYear = int(curDayStr[4:8])
					curMonth = int(curDayStr[8:10])
					curDay = int(curDayStr[10:12])
					curDate = datetime.datetime(curYear, curMonth, curDay, 0, 0, 0)
				
					if curDate not in activeDateTime:
						os.system('rmdir /s /q "%s/%d/%s"'%(baseDir, year, curDirName))
						print 'deleted %s/%d/%s'%(baseDir, year, curDirName)
			
			# or if there are just .tar.z files
			for curFileName in filenames:
				fileNameParts = curFileName.split('.')
				if fileNameParts[1] == 'tar' and fileNameParts[2] == 'Z':
					curYear = int(fileNameParts[0][4:8])
					curMonth = int(fileNameParts[0][8:10])
					curDay = int(fileNameParts[0][10:12])
					curDate = datetime.datetime(curYear, curMonth, curDay, 0, 0, 0)
					
					if curDate not in activeDateTime:
						filePath = '"%s/%d/%s"'%(baseDir, year, curFileName)
						filePath = filePath.replace('/', '\\')
						os.system('del %s'%filePath)
						print 'deleted %s'%filePath
					
			break
	return activeFolders

def processNewRadarData(years):	
	folders = removeInactiveRadarData(years)
	organizeRadarTars(years)
	radarToRv3(folders)
	unzipRadarData(years = years, clean = True)
